#!/usr/bin/python

"""
This program draws MUMMER alignment results and shows connecting segments
based on % identity.

Usage: dissertation_DrawMUMMER.py g1.gbk g2.gbk g1.ptt g2.ptt cogs.t.list mummer.coord 
Examples:
Go to this directory:
/host/Users/JS/UH-work/gloeobacter/final_work/comparisons

dissertation_DrawMUMMERwithPtt.py NC_007776.gbk NC_007775.gbk NC_007776.ptt 
    NC_007775.ptt orthologs/orthomcl/cogs.t.list NC_007776_vs_NC_007775.coords
dissertation_DrawMUMMERwithPtt.py NC_007516.gbk NC_007513.gbk NC_007516.ptt 
    NC_007513.ptt orthologs/orthomcl/cogs.t.list NC_007516_vs_NC_007513.coords
dissertation_DrawMUMMERwithPtt.py NC_011748.gbk NC_008253.gbk NC_011748.ptt 
    NC_008253.ptt orthologs/orthomcl/cogs.t.list NC_011748_vs_NC_008253.coords
dissertation_DrawMUMMERwithPtt.py NC_002737.gbk NC_007297.gbk NC_002737.ptt 
    NC_007297.ptt orthologs/orthomcl/cogs.t.list NC_002737_vs_NC_007297.coords

Steps:
1. Download fna, gbk, and ptt files
2. Run MUMMER
3. Run this program

Author: Jimmy Saw
Date of last update: 04-23-2012

"""

import sys
import re
import matplotlib.pyplot as plt
import pylab
import matplotlib
from matplotlib import mpl
from matplotlib.patches import Rectangle
from matplotlib.transforms import Bbox
from Bio import SeqIO
from Bio.SeqUtils import GC
import matplotlib.patches as mpatch

#Regex and other stuffs
cogcat = re.compile('\[(.*)\]\t(\w+)\t.*')
#pttcog = re.compile('(COG\d+)(\w).*')
pttcog = re.compile('(COG\d{4})(\w).*')

cogdict = {
    'J' : '#2B60DE',
    'A' : '#F6358A',
    'K' : '#B048B5',
    'L' : '#8E35EF',
    'B' : '#D16587',
    'D' : '#C38EC7',
    'Y' : '#52F3FF',
    'V' : '#3EA99F',
    'T' : '#254117',
    'M' : '#41A317',
    'N' : '#00FF00',
    'Z' : '#FFFF00',
    'W' : '#FDD017',
    'U' : '#F88017',
    'O' : '#F660AB',
    'C' : '#FF0000',
    'G' : '#FAAFBA',
    'E' : '#7F5A58',
    'F' : '#C8B560',
    'H' : '#D2B9D3',
    'I' : '#C12869',
    'P' : '#57E964',
    'Q' : '#BCE954',
    'R' : '#F87431',
    'S' : '#ADA96E',
}



##Genome 1 Genbank file
g1seq = SeqIO.read(sys.argv[1], "gb")
g1length = len(g1seq.seq)
g1featdict = {}
for feat in g1seq.features:
    if feat.type == 'CDS':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g1featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 2 Genbank file
g2seq = SeqIO.read(sys.argv[2], "gb")
g2length = len(g2seq.seq)
g2featdict = {}
for feat in g2seq.features:
    if feat.type == 'CDS':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'tRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat
    if feat.type == 'rRNA':
        g2featdict[feat.qualifiers['locus_tag'][0]] = feat

##Genome 1 ptt file
g1cogs = {}
g1pttfile = open(sys.argv[3], "rU")
g1ptt = g1pttfile.readlines()
for line in g1ptt[3:]:
    c = line.split('\t')
    g1ltag = c[5]
    g1gene = c[4]
    #g1cogcat = '-'
    g1cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g1cogcat = p.group(2)
        g1cog = p.group(1)
    #g1cogs[g1ltag] = g1cogcat
    g1cogs[g1ltag] = g1cog

##Genome 2 ptt file
g2cogs = {}
g2pttfile = open(sys.argv[4], "rU")
g2ptt = g2pttfile.readlines()
for line in g2ptt[3:]:
    c = line.split('\t')
    g2ltag = c[5]
    g2gene = c[4]
    #g2cogcat = '-'
    g2cog = '-'
    if pttcog.match(c[7]):
        p = pttcog.match(c[7])
        #g2cogcat = p.group(2)
        g2cog = p.group(1)
    #g2cogs[g2ltag] = g2cogcat
    g2cogs[g2ltag] = g2cog

cogcatfile = open(sys.argv[5], "rU")
cfl = cogcatfile.readlines()

cogcatdict = {}

for line in cfl:
    tmp = line.strip()
    if cogcat.match(tmp):
        pattern = cogcat.match(tmp)
        cogcatdict[pattern.group(2)] = pattern.group(1)[0]

cogcatfile.close()

#spanx1 = int(sys.argv[7])
#spanx2 = int(sys.argv[8])

glist = []

genome1x = [0, g1length]
genome1y = [2, 2]
genome2x = [0, g2length]
genome2y = [12, 12]

largergenome = 0

if g1length > g2length:
    largergenome = g1length
else:
    largergenome = g2length

##Start plotting
fig = plt.figure(1, figsize=(14,4))
#ax1 = fig.add_subplot(211) #makes the subplot and squeezes the figure to half panel
ax1 = fig.add_subplot(111) #makes the full figure plot. larger.
ax1.plot(genome1x, genome1y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)
ax1.plot(genome2x, genome2y, color='#FFFFFF', marker='|', markersize=8.0,
    mec='black', ls='-', lw=2.0)

ax1.axis([0, largergenome, 0, 14])
#ax1.axis([spanx1, spanx2, 0, 14])

for k, v in g1featdict.iteritems():
    g1feat = v
    g1locustag = g1feat.qualifiers['locus_tag'][0]
    g1start = g1feat.location._start.position
    g1stop = g1feat.location._end.position
    g1size = g1stop - g1start + 1
    g1mid = g1start + ((g1stop - g1start) / 2.0)
    g1desc = g1feat.qualifiers['product'][0]
    g1gene = ""
    if g1feat.qualifiers.has_key('gene'):
        g1gene = g1feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g1gene = g1feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g1locustag in g1cogs:
        if g1cogs[g1locustag] != '-':
            #cogcolor = cogdict[g1cogs[g1locustag]]
            cogcolor = cogdict[cogcatdict[g1cogs[g1locustag]]]
    if g1feat.type == 'tRNA':
        cogcolor = '#800000'
    if g1feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g1feat.strand == -1:
        rect = Rectangle((g1start, 4.0), g1size, 0.5, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g1mid, 4.5, g1gene, fontsize=8, color='black', rotation=45)
    else:
        rect = Rectangle((g1start, 4.5), g1size, 0.5, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g1mid, 5.5, g1gene, fontsize=8, color='black', rotation=45)

for k, v in g2featdict.iteritems():
    g2feat = v
    g2locustag = g2feat.qualifiers['locus_tag'][0]
    g2start = g2feat.location._start.position
    g2stop = g2feat.location._end.position
    g2size = g2stop - g2start + 1
    g2mid = g2start + ((g2stop - g2start) / 2.0)
    g2desc = g2feat.qualifiers['product'][0]
    g2gene = ""
    if g2feat.qualifiers.has_key('gene'):
        g2gene = g2feat.qualifiers['gene'][0] #displays gene name
        #g1gene = g1desc #displays product description
    else:
        g2gene = g2feat.qualifiers['locus_tag'][0] #displays locus tag
        #g1gene = g1desc #displays product description
    cogcolor = '#D3D3D3' #base color
    if g2locustag in g2cogs:
        if g2cogs[g2locustag] != '-':
            #cogcolor = cogdict[g2cogs[g2locustag]]
            cogcolor = cogdict[cogcatdict[g2cogs[g2locustag]]]
    if g2feat.type == 'tRNA':
        cogcolor = '#800000'
    if g2feat.type == 'rRNA':
        cogcolor = '#9400D3'
        g1gene = g1desc
    if g2feat.strand == -1:
        rect = Rectangle((g2start, 9.0), g2size, 0.5, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g2mid, 9.5, g2gene, fontsize=8, color='black', rotation=45)
    else:
        rect = Rectangle((g2start, 9.5), g2size, 0.5, fc=cogcolor, 
            ec=cogcolor, alpha=0.5)
        plt.gca().add_patch(rect)
        #ax1.text(g2mid, 10.5, g2gene, fontsize=8, color='black', rotation=45)

#ax1.text(len(g1seq.seq)+10000, 4.2, "-", fontsize=10, color='black')
#ax1.text(len(g1seq.seq)+10000, 4.6, "+", fontsize=10, color='black')
#ax1.text(len(g2seq.seq)+10000, 9.2, "-", fontsize=10, color='black')
#ax1.text(len(g2seq.seq)+10000, 9.6, "+", fontsize=10, color='black')
ax1.annotate(g1seq.annotations['organism'], xy=(0.5, 0.1), 
    xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)
ax1.annotate(g2seq.annotations['organism'], xy=(0.5, 0.9), 
    xycoords='axes fraction', horizontalalignment='center', verticalalignment='center', fontsize=10)

##Parse MUMMER alignment file

mummerfile = open(sys.argv[6], "rU")
mfl = mummerfile.readlines()
for line in mfl[4:]:
    c = line.split('\t')
    g1start = int(c[0])
    g1stop = int(c[1])
    g2start = int(c[2])
    g2stop = int(c[3])
    ident = float(c[6])
    fillcolor = '#AAAAAA'
    if ident >= 90:
        fillcolor = '#FF0000'
    elif ident >= 80:
        fillcolor = '#71C671'
    elif ident >= 70:
        fillcolor = '#7171C6'
    elif ident >= 60:
        fillcolor = '#CDB5CD'
    else:
        fillcolor = '#C5C1AA'
    x = [g1start, g2start, g2stop, g1stop]
    y = [5, 9, 9, 5]
    ax1.fill(x, y, color=fillcolor, alpha=0.2)

"""
#This segment is an example of how to draw polygon connecting genomic regions
a = [2015000, 2018000, 2030000, 2020000]
a1 = [2015000, 2030000, 2018000, 2020000]
b = [6, 7, 7, 6]
ax1.fill(a1, b, 'r', alpha=0.2)
"""


frame1 = plt.gca()
for tick in frame1.axes.get_yticklines():
    tick.set_visible(False)
for y in frame1.axes.get_yticklabels():
    y.set_visible(False)
ax1.grid(False)

plt.show()
